import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim

from defense.base_defense import BaseDefense

class Defense(BaseDefense):
    def __init__(self, model_name, config):
        super().__init__('PAT', model_name, config, retrain = True)

    #进行PGD攻击
    def PGD_attack(self, images=None, labels=None):
        self.model.eval()
        natural_images = images.cpu().numpy()

        copy_images = natural_images.copy()
        copy_images = copy_images + np.random.uniform(-self.config['eps'], self.config['eps'], copy_images.shape).astype('float32')

        for i in range(self.config['step_num']):
            v_copy_images = torch.from_numpy(copy_images).to(self.model.device)
            v_copy_images.requires_grad = True

            preds = self.model(v_copy_images)
            loss = F.cross_entropy(preds, labels)
            gradient = torch.autograd.grad(loss, v_copy_images)[0]
            sign = torch.sign(gradient).cpu().numpy()

            copy_images = copy_images + self.config['step_size'] * sign

            copy_images = np.clip(copy_images, natural_images - self.config['eps'], natural_images + self.config['eps'])
            copy_images = np.clip(copy_images, 0.0, 1.0)

        return torch.from_numpy(copy_images).to(self.model.device)

    # 单轮防御训练
    def train(self, loader, opt, epoch):

        for ind, (images, labels) in enumerate(loader):
            images = images.to(self.model.device)
            labels = labels.to(self.model.device)

            # 进行PGD攻击
            advs = self.PGD_attack(images, labels) # advs:对抗图像

            self.model.train()

            output = self.model(images)
            loss = F.cross_entropy(output, labels)  # 原始图像损失

            adv_output = self.model(advs)
            adv_loss = F.cross_entropy(adv_output, labels)  # 混淆图像损失

            total_loss = 0.5 * (loss + adv_loss)  # 总损失

            opt.zero_grad()
            total_loss.backward()
            opt.step()

            ind += 1
            print('\rEpoch {}: [{}/{} ({:.0f}%)] \tLoss: {:.4f} AdvLoss: {:.4f} TotalLoss: {:.4f}'.format(
                epoch, (ind + 1) * len(images), len(loader) * len(images),
                (ind + 1) / len(loader) * 100, loss, adv_loss, total_loss), end=' ')

        print()

    def defend(self):
        batch_size = self.model.training_params['batch_size']
        # 加载训练数据
        train_loader, valid_loader = self.model.get_train_valid_loader(batch_size=batch_size,
                                                                       valid_size=0.1,
                                                                       transform_train=True,
                                                                       shuffle=True)
        # 使用SGD优化函数
        opt = self.model.get_optimizer()

        best_ratio = 0
        best_state_dict = None
        for epoch in range(self.model.training_params['num_epochs']):
            self.train(train_loader, opt, epoch)
            ratio = self.model.test(valid_loader)
            self.model.adjust_optimizer(opt, epoch)
            if ratio > best_ratio:
                best_ratio = ratio
                best_state_dict = self.model.state_dict_clone()  # 保存此时网络参数
                print('Validate: Accuracy increased')
            else:
                print('Validate: Accuracy decreased')
            print()
        return best_state_dict  # 返回最优参数